package com.chengxiaoxiao.ai_demo.service.impl;

import cn.hutool.core.collection.CollUtil;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.chengxiaoxiao.ai_demo.config.KnowledgeThreadLocal;
import com.chengxiaoxiao.ai_demo.entity.ConversationInfo;
import com.chengxiaoxiao.ai_demo.exception.GlobalException;
import com.chengxiaoxiao.ai_demo.mapper.ConversationInfoMapper;
import com.chengxiaoxiao.ai_demo.service.ConversationInfoService;
import com.chengxiaoxiao.ai_demo.service.IndexInfoService;
import com.chengxiaoxiao.ai_demo.service.InstructionInfoService;
import com.chengxiaoxiao.ai_demo.service.OriginalMessageService;
import com.chengxiaoxiao.ai_demo.service.ai.IndexAiService;
import com.chengxiaoxiao.ai_demo.vo.CodeMsg;
import com.chengxiaoxiao.ai_demo.vo.form.ChatForm;
import com.chengxiaoxiao.ai_demo.vo.form.ContextForm;
import com.chengxiaoxiao.ai_demo.vo.vo.ChatResponseVo;
import com.chengxiaoxiao.ai_demo.vo.vo.IndexTreeVo;
import dev.langchain4j.service.TokenStream;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Sinks;

import java.util.List;
import java.util.Objects;

/**
 * 会话信息实现类
 *
 * @Description:
 * @Author Cheng Xiaoxiao
 * @Date: 2025/3/12 15:56
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class ConversationInfoServiceImpl extends ServiceImpl<ConversationInfoMapper, ConversationInfo> implements ConversationInfoService {
    final ConversationInfoMapper conversationInfoMapper;
    final IndexInfoService indexInfoService;
    final IndexAiService indexAiService;
    final InstructionInfoService instructionInfoService;
    final OriginalMessageService originalMessageService;

    /**
     * 创建新会话
     *
     * @return 创建的会话信息
     */
    @Override
    public ConversationInfo create() {
        ConversationInfo conversationInfo = new ConversationInfo().setTitle("新会话");
        conversationInfoMapper.insert(conversationInfo);
        return conversationInfo;
    }

    /**
     * 根据ID修改会话信息
     *
     * @param id               会话ID
     * @param conversationInfo 会话信息
     */
    @Override
    public ConversationInfo editById(String id, ConversationInfo conversationInfo) {
        //如果会话不存在，错误提示
        if (conversationInfoMapper.selectById(id) == null) {
            throw new GlobalException(CodeMsg.CONVERSATION_NOT_FOUND);
        }

        conversationInfo.setId(id);
        conversationInfoMapper.updateById(conversationInfo);

        //再次查询最新信息
        return conversationInfoMapper.selectById(id);
    }

    /**
     * 删除会话
     *
     * @param id 会话ID
     */
    @Override
    public void deleteById(String id) {
        //如果会话不存在，错误提示
        if (conversationInfoMapper.selectById(id) == null) {
            throw new GlobalException(CodeMsg.CONVERSATION_NOT_FOUND);
        }

        conversationInfoMapper.deleteById(id);
    }


    /**
     * 将 IndexTreeVo 列表转换为```
     *
     * @param indexTreeVos 指标信息列表
     * @return markdown字符串
     */
    private String convertToMarkdown(List<IndexTreeVo> indexTreeVos) {
        StringBuilder markdown = new StringBuilder();
        for (IndexTreeVo vo : indexTreeVos) {
            convertToMarkdownRecursive(vo, markdown, 0);
        }
        return markdown.toString();
    }

    /**
     * 递归将 IndexTreeVo 对象转换为```
     *
     * @param vo          指标信息对象
     * @param markdown```
     * @param level       当前层级
     */
    private void convertToMarkdownRecursive(IndexTreeVo vo, StringBuilder markdown, int level) {
        // 添加标题
        markdown.append("#".repeat(level + 1)).append(" ").append(vo.getName()).append("\n");
        // 添加描述
        markdown.append(vo.getDescription()).append("\n");
        // 递归处理子节点
        if (CollUtil.isNotEmpty(vo.getChildren())) {
            for (IndexTreeVo child : vo.getChildren()) {
                convertToMarkdownRecursive(child, markdown, level + 1);
            }
        }
    }

    /**
     * 根据会话进行聊天
     *
     * @param id       会话ID
     * @param chatForm 会话表单
     * @return 回答的结果
     */
    @Override
    public Flux<ChatResponseVo> chatStream(String id, ChatForm chatForm) {
        //存储原始用户信息
        originalMessageService.saveUserMessage(id, chatForm);

        //处理上下文信息
        String contextMessage = formatContext(chatForm.getContext());

        //如果存在指令，则进行替换为具体的提示词
        contextMessage = contextMessage + "\n" +
                "<user_query>\n" + instructionInfoService.buildMessage(chatForm.getText()) + "\n</user_query>\n";

        //保存知识库信息，供后续使用
        KnowledgeThreadLocal.setKnowledgeBaseId(chatForm.getKnowledgeBaseId());

        TokenStream tokenStream = indexAiService.chatStream(id, contextMessage);

        //记录数据
        StringBuilder stringBuilder = new StringBuilder();

        //转换为固定格式的流数据
        Sinks.Many<ChatResponseVo> sink = Sinks.many().unicast().onBackpressureBuffer();
        Objects.requireNonNull(sink);
        TokenStream var10000 = tokenStream
                .onPartialResponse(s -> {
                    sink.tryEmitNext(new ChatResponseVo().setData(s));
                    stringBuilder.append(s);
                })
                .onCompleteResponse((ignored) -> {
                    sink.tryEmitComplete();
                    originalMessageService.saveAiMessage(id, stringBuilder.toString());
                });
        Objects.requireNonNull(sink);
        var10000.onError(sink::tryEmitError).start();
        return sink.asFlux();
    }

    /**
     * 构建上下文
     *
     * @param context 上下文信息
     * @return 上下文markdown
     */
    private String formatContext(List<ContextForm> context) {
        StringBuilder contextBuilder = new StringBuilder();
        //如果有选中上下文
        if (CollUtil.isNotEmpty(context)) {
            contextBuilder.append("<context_selection>\n");
            for (ContextForm contextForm : context) {
                //获取指标信息
                List<IndexTreeVo> indexTreeVos = indexInfoService.getIndexInfo(contextForm.getId());
                //将指标信息拼接为有层级的信息
                contextBuilder.append(convertToMarkdown(indexTreeVos));
            }
            contextBuilder.append("</context_selection>\n");
        }
        return contextBuilder.toString();
    }
}